import torch
import numpy as np
import matplotlib.pyplot as plt
import cv2


def show_image(image, title=''):
    # image is [H, W, 3]
    plt.imshow(image)
    plt.title(title, fontsize=16)
    plt.axis('off')
    return


def run_one_image(img, model, mask_ratio=0.75):
    model = model.cuda()
    x = torch.tensor(img)

    # make it a batch-like
    x = x.unsqueeze(dim=0)
    x = torch.einsum('nhwc->nchw', x)
    x = x.cuda()

    # run MAE
    loss, y, mask = model(x.float(), mask_ratio)
    y = model.unpatchify(y)
    y = torch.einsum('nchw->nhwc', y).detach().cpu()

    # visualize the mask
    mask = mask.detach()
    mask = mask.unsqueeze(-1).repeat(1, 1, model.patch_embed.patch_size[0] ** 2 * 3)  # (N, H*W, p*p*3)
    mask = model.unpatchify(mask)  # 1 is removing, 0 is keeping
    mask = torch.einsum('nchw->nhwc', mask).detach().cpu()

    x = torch.einsum('nchw->nhwc', x).cpu()

    # masked image
    im_masked = x * (1 - mask)

    # MAE reconstruction pasted with visible patches
    im_paste = x * (1 - mask) + y * mask

    # make the plt figure larger
    plt.rcParams['figure.figsize'] = [24, 24]

    plt.subplot(1, 4, 1)
    show_image(x[0], "original")

    plt.subplot(1, 4, 2)
    show_image(im_masked[0], "masked")

    plt.subplot(1, 4, 3)
    show_image(y[0], "reconstruction")

    plt.subplot(1, 4, 4)
    show_image(im_paste[0], "reconstruction + visible")
    plt.savefig('test_one_image.jpg')
    plt.show()
    # pylab.show()


if __name__ == '__main__':
    import datetime
    print(datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S") + ' run ' + __file__.split('/')[-1])
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument('--model_path', action='store', type=str, default='mae_transistor.pth')
    parser.add_argument('--image_path', action='store', type=str, default='/kaggle/input/mvtecad-mvtec-anomaly-detection/mvtec_anomaly_detection/transistor/test/cut_lead/001.png')
    parser.add_argument('--mask_ratio', action='store', type=float, default=0.8)
    parser.add_argument('--seed', action='store', type=int, default=24)
    args = parser.parse_args()

    model_mae = torch.load(args.model_path)
    print('Model loaded.')

    image = cv2.imread(args.image_path)
    image = cv2.resize(image, dsize=(224, 224))
    image = np.array(image).reshape((image.shape[0], image.shape[1], image.shape[2])).astype(np.float32) / 255.0

    # make random mask reproducible
    torch.manual_seed(args.seed)
    print('MAE with pixel reconstruction:')
    run_one_image(image, model_mae, args.mask_ratio)
